// Copyright (C) 2018-2023 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "transformations/resolve_names_collisions.hpp"

#include <algorithm>
#include <memory>

#include "itt.hpp"
#include "openvino/op/util/multi_subgraph_base.hpp"

namespace {

void collect_name_collisions_map(const std::shared_ptr<ov::Model>& model,
                                 std::unordered_map<std::string, std::list<ov::Node*>>& name_collisions_map) {
    for (const auto& node : model->get_ordered_ops()) {
        // Collect a names collision map for all nodes in the graph
        const auto& friendly_name = node->get_friendly_name();
        name_collisions_map[friendly_name].emplace_back(node.get());
        if (auto msn = ov::as_type_ptr<ov::op::util::MultiSubGraphOp>(node)) {
            for (const auto& body : msn->get_functions()) {
                collect_name_collisions_map(body, name_collisions_map);
            }
        }
    }
}

}  // namespace

bool ov::pass::ResolveNameCollisions::run_on_model(const std::shared_ptr<ov::Model>& model) {
    // Next containers are used to fix collisions in autogenerated names
    // The final list of nodes with collisions
    RUN_ON_MODEL_SCOPE(ResolveNameCollisions);
    std::vector<Node*> nodes_with_conflicts;
    std::unordered_map<std::string, std::list<Node*>> visited_nodes;

    collect_name_collisions_map(model, visited_nodes);

    for (const auto& it : visited_nodes) {
        const auto& same_named_ops = it.second;
        if (same_named_ops.size() < 2) {
            continue;
        }

        auto update_name = [&](ov::Node* op, int64_t& cnt) {
            // add a prefix "_counter" to the autogenerated name to make it unique.
            const auto& old_name = op->get_friendly_name();
            auto new_name = old_name + "_" + std::to_string(cnt++);
            while (visited_nodes.find(new_name) != visited_nodes.end()) {
                new_name = old_name + "_" + std::to_string(cnt++);
            }
            op->set_friendly_name(new_name);
        };

        if (m_resolve_all_names) {
            int64_t cnt = 0;
            for (const auto& op : same_named_ops) {
                // we don't need to change name of the first op in this list.
                // the expected result: Name, Name_1, Name_2 etc
                if (cnt == 0) {
                    cnt++;
                    continue;
                }
                update_name(op, cnt);
            }
        } else {
            int64_t cnt = 2;
            for (const auto& op : same_named_ops) {
                // check if op has OV autogenerated friendly name.
                // unique and friendly names have to be equal for the autogenerated name.
                bool is_autogenerated = op->m_friendly_name.empty();
                if (!is_autogenerated) {
                    continue;
                }
                update_name(op, cnt);
            }
        }
    }
    return false;
}
